from langchain.tools import BaseTool
from langchain.chains import create_sql_query_chain
from langchain_community.utilities import SQLDatabase
from pydantic import BaseModel, Field, ConfigDict
from typing import Any, Type, Optional
from dotenv import load_dotenv
from model.my_chat_model import ChatModel
import os


# 定义输入参数的数据模型类
class LoginInput(BaseModel):
    account: str = Field(..., description="用户名")
    password: str = Field(..., description="密码")


# 定义工具类
class LoginTool(BaseTool):
    # 定义模型，是否允许输入参数，输入参数的数据模型类
    model_config = ConfigDict(arbitrary_types_allowed=True)

    def __init__(self, **kwargs: Any):
        super().__init__(
            name="get_login_tool",
            description="主要用于完成系统的登录功能，必须输入的参数是用户名和密码",
            **kwargs
        )

    # 定义工具参数
    args_schema: Type[BaseModel] = LoginInput

    # 定义工具方法
    def _run(self, account: str, password: str):
        # 1 获取大模型
        chat = ChatModel()
        llm = chat.get_line_model()
        # 2 创建数据库链接
        db = SQLDatabase.from_uri(
            "mysql+pymysql://root:root@localhost:3306/ai",
            include_tables=["account_password"],
        )
        # 3 创建链
        chain = create_sql_query_chain(llm, db)
        # 4 提问
        question = f"请根据条件用户名是{account}和密码是{password}查询用户信息"
        sql = chain.invoke({"question": question})
        print(sql)

        if "```sql" in sql:
            sql = sql.split("```sql")[1].split("```")[0]
        print(sql)
        # sql = "SELECT `user_id`, `user_name`, `email` FROM `user_info` WHERE `user_name` = '张三' AND `user_pwd` = '1234' LIMIT 5;"
        rs = db.run(sql)
        print(rs)
        return rs
